class TaskBuffer:
    def __init__(self, max_size=100):#, initial_storage=0):
        # super.__init__(max_size)
        self.storage = list()
        self.max_size = max_size

    def add(self, data):
        self.storage.append(data)
        if len(self.storage) > self.max_size:
            self.storage = self.storage[1:]
        else:
            pass

    def get_buffer(self):
        return self.storage[::-1]

    def get_last_obj(self):
        if self.storage:
            return self.storage[-1]
        else:
            return None

# For TD3 agents
class ReplayBuffer(object):
	def __init__(self, max_size=1e6):
		self.storage = []
		self.max_size = max_size
		self.ptr = 0

	def add(self, data):
		if len(self.storage) == self.max_size:
			self.storage[int(self.ptr)] = data
			self.ptr = (self.ptr + 1) % self.max_size
		else:
			self.storage.append(data)

	def sample(self, batch_size):
		ind = np.random.randint(0, len(self.storage), size=batch_size)
		x, y, u, r, d = [], [], [], [], []

		for i in ind:
			X, Y, U, R, D = self.storage[i]
			x.append(np.array(X, copy=False))
			y.append(np.array(Y, copy=False))
			u.append(np.array(U, copy=False))
			r.append(np.array(R, copy=False))
			d.append(np.array(D, copy=False))

		return np.array(x), np.array(y), np.array(u), np.array(r).reshape(-1, 1), np.array(d).reshape(-1, 1)
